<!--
@license
Copyright 2016 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->

<link rel="import" href="../iron-flex-layout/iron-flex-layout.html">
<link rel="import" href="../iron-icons/iron-icons.html">
<link rel="import" href="../paper-button/paper-button.html">
<link rel="import" href="../paper-input/paper-input.html">
<link rel="import" href="../paper-toggle-button/paper-toggle-button.html">
<link rel="import" href="../polymer/polymer.html">
<link rel="import" href="../tf-graph-common/tf-graph-common.html">
<link rel="import" href="tf-graph-scene.html">

<dom-module id="tf-graph">
<template>
<style>
.container {
  width: 100%;
  height: 100%;
  background: white;
  box-shadow: 0 1px 5px rgba(0,0,0,0.2);
}

.vertical {
  width:100%;
  height:100%;
  @apply --layout-vertical;
}

.auto {
  @apply --layout-flex-auto;
  @apply --layout-vertical;
}

h2 {
  text-align: center;
}

paper-button {
  text-transform: none;
}
</style>
<div class="container">
  <div class="vertical">
    <template is="dom-if" if="[[title]]">
      <h2>[[title]]</h2>
    </template>
    <tf-graph-scene id="scene" class="auto"
          render-hierarchy="[[renderHierarchy]]"
          highlighted-node="[[_getVisible(highlightedNode)]]"
          selected-node="{{selectedNode}}"
          selected-edge="{{selectedEdge}}"
          color-by="[[colorBy]]"
          progress="[[progress]]"
          node-context-menu-items="[[nodeContextMenuItems]]"
          node-names-to-health-pills="[[nodeNamesToHealthPills]]"
          health-pill-step-index="{{healthPillStepIndex}}"
          handle-edge-selected="[[handleEdgeSelected]]"
    ></tf-graph-scene>
  </div>
</div>
</template>
</dom-module>

<script>
Polymer({

  is: 'tf-graph',

  properties: {
    graphHierarchy: {
      type: Object,
      notify: true,
      observer: '_graphChanged'
    },
    basicGraph: Object,
    stats: Object,
    devicesForStats: Object,
    hierarchyParams: Object,
    progress: {
      type: Object,
      notify: true,
    },
    title: String,
    selectedNode: {
      type: String,
      notify: true,
    },
    // An EdgeData object if an edge is selected. Otherwise, null.
    selectedEdge: {
      type: Object,
      notify: true,
    },
    // The <g> element of the last selected edge. Used to mark the edge as selected.
    _lastSelectedEdgeGroup: Object,
    highlightedNode: {
      type: String,
      notify: true
    },
    /** What to color the nodes by (compute time, memory, device etc.) */
    colorBy: String,
    colorByParams: {
      type: Object,
      notify: true,
      readOnly: true, // Produces and doesn't consume.
    },
    renderHierarchy: {
      type: Object,
      readOnly: true,
      notify: true,
    },
    // An array of ContextMenuItem objects. Items that appear in the context
    // menu for a node.
    nodeContextMenuItems: Array,
    _renderDepth: {
      type: Number,
      value: 1
    },
    _allowGraphSelect: {
      type: Boolean,
      value: true
    },
    // A mapping between node name to the tf.graph.scene.HealthPill to render.
    nodeNamesToHealthPills: Object,
    // The step of health pills to show throughout the graph.
    healthPillStepIndex: Number,
    /**
     * A function with signature EdgeThicknessFunction that computes the
     * thickness of a given edge.
     *
     * We initialize with the empty string value so that the observer that
     * builds the graph hierarchy is called even if this is not defined.
     */
    edgeWidthFunction: {
      type: Object,
      value: '',
    },
    /**
     * An optional function that takes a node selected event (whose `detail`
     * property is the selected node ... which could be null if a node is
     * deselected). Called whenever a node is selected or deselected.
     *
     * We initialize with the empty string value so that the observer that
     * builds the graph hierarchy is called even if this is not defined.
     * @type {Function}
     */
    handleNodeSelected: {
      type: Object,
      value: '',
    },
    /**
     * An optional function that computes the label for an edge. Should
     * implement the EdgeLabelFunction signature.
     *
     * We initialize with the empty string value so that the observer that
     * builds the graph hierarchy is called even if this is not defined.
     * @type {Function}
     */
    edgeLabelFunction: {
      type: Object,
      value: '',
    },
    /**
     * An optional callback that implements the
     * tf.graph.edge.EdgeSelectionCallback signature. If provided, edges are
     * selectable, and this callback is run when an edge is selected.
     *
     * We initialize with the empty string value so that the observer that
     * builds the graph hierarchy is called even if this is not defined.
     * @type {Function}
     */
    handleEdgeSelected: {
      type: Object,
      value: '',
    },
  },
  observers: [
    '_statsChanged(stats, devicesForStats)',
    // We must re-render the graph hierarchy if any handlers are defined.
    // Otherwise, the graph hierarchy might be created before the handlers are
    // set on the hierarchy, and the handlers will not be taken into
    // consideration by graph rendering logic. That would unfortunately for
    // instance result in the edge width function sometimes failing to take
    // effect on some page loads.
    '_buildNewRenderHierarchy(graphHierarchy, edgeWidthFunction, handleNodeSelected, edgeLabelFunction, handleEdgeSelected)',
    '_selectedNodeChanged(selectedNode)',
    '_selectedEdgeChanged(selectedEdge)',
  ],
  /**
   * Pans to a node. Assumes that the node exists.
   * @param nodeName {string} The name of the node to pan to.
   */
  panToNode(nodeName) {
    this.$$('tf-graph-scene').panToNode(nodeName);
  },
  _buildNewRenderHierarchy(graphHierarchy) {
    this._buildRenderHierarchy(graphHierarchy);
  },
  _statsChanged: function(stats, devicesForStats) {
    if (this.graphHierarchy) {
      if (stats && devicesForStats) {
        tf.graph.joinStatsInfoWithGraph(this.basicGraph, stats, devicesForStats);
        tf.graph.hierarchy.joinAndAggregateStats(this.graphHierarchy, stats);
      }
      // Recompute the rendering information.
      this._buildRenderHierarchy(this.graphHierarchy);
    }
  },
  _buildRenderHierarchy: function(graphHierarchy) {
    tf.graph.util.time('new tf.graph.render.Hierarchy', function() {
      if (graphHierarchy.root.type !== tf.graph.NodeType.META) {
        // root must be metanode but sometimes Polymer's dom-if has not
        // remove tf-graph element yet in <tf-node-info>
        // and thus mistakenly pass non-metanode to this module.
        return;
      }
      var renderGraph = new tf.graph.render.RenderGraphInfo(
          graphHierarchy, !!this.stats /** displayingStats */);
      renderGraph.edgeLabelFunction = this.edgeLabelFunction;
      renderGraph.edgeWidthFunction = this.edgeWidthFunction;

      // Producing the 'color by' parameters to be consumed
      // by the tf-graph-controls panel. It contains information about the
      // min and max values and their respective colors, as well as list
      // of devices with their respective colors.
      function getColorParamsFromScale(scale) {
        return {
          minValue: scale.domain()[0],
          maxValue: scale.domain()[1],
          startColor: scale.range()[0],
          endColor: scale.range()[1]
        };
      }

      this._setColorByParams({
        compute_time: getColorParamsFromScale(renderGraph.computeTimeScale),
        memory: getColorParamsFromScale(renderGraph.memoryUsageScale),
        device: _.map(renderGraph.deviceColorMap.domain(),
            function(deviceName) {
          return {
            device: deviceName,
            color: renderGraph.deviceColorMap(deviceName)
          };
        }),
        xla_cluster: _.map(renderGraph.xlaClusterColorMap.domain(),
            function(xlaClusterName) {
          return {
            xla_cluster: xlaClusterName,
            color: renderGraph.xlaClusterColorMap(xlaClusterName)
          };
        }),
      });
      this._setRenderHierarchy(renderGraph);
      this.async(function() {
        this.fire("rendered");
      });
    }.bind(this));
  },
  _getVisible: function(name) {
    if (!name) {
      return name;
    }
    return this.renderHierarchy.getNearestVisibleAncestor(name);
  },
  listeners: {
    'graph-select': '_graphSelected',
    'disable-click': '_disableClick',
    'enable-click': '_enableClick',
    // Nodes
    'node-toggle-expand': '_nodeToggleExpand',
    'node-select': '_nodeSelected',
    'node-highlight': '_nodeHighlighted',
    'node-unhighlight': '_nodeUnhighlighted',
    'node-toggle-extract': '_nodeToggleExtract',
    'node-toggle-seriesgroup': '_nodeToggleSeriesGroup',
    // Edges
    'edge-select': '_edgeSelected',

    // Annotations

    /* Note: currently highlighting/selecting annotation node has the same
      * behavior as highlighting/selecting actual node so we point to the same
      * set of event listeners.  However, we might redesign this to be a bit
      * different.
      */
    'annotation-select': '_nodeSelected',
    'annotation-highlight': '_nodeHighlighted',
    'annotation-unhighlight': '_nodeUnhighlighted',
  },
  _graphChanged: function() {
    // When a new graph is loaded, fire this event so that there is no
    // info-card being displayed for the previously-loaded graph.
    this.fire('graph-select');
  },
  _graphSelected: function(event) {
    // Graph selection is not allowed during an active zoom event, as the
    // click seen during a zoom/pan is part of the zooming and does not
    // indicate a user desire to click on a specific section of the graph.
    if (this._allowGraphSelect) {
      this.set('selectedNode', null);
      this.set('selectedEdge', null);
    }
    // Reset this variable as a bug in d3 zoom behavior can cause zoomend
    // callback not to be called if a right-click happens during a zoom event.
    this._allowGraphSelect = true;
  },
  _disableClick: function(event) {
    this._allowGraphSelect = false;
  },
  _enableClick: function(event) {
    this._allowGraphSelect = true;
  },
  // Called when the selected node changes, ie there is a new selected node or
  // the current one is unselected.
  _selectedNodeChanged(selectedNode) {
    if (this.handleNodeSelected) {
      // A higher-level component provided a callback. Run it.
      this.handleNodeSelected(selectedNode);
    }
  },
  // Called when the selected edge changes, ie there is a new selected edge or
  // the current one is unselected.
  _selectedEdgeChanged(selectedEdge) {
    this._deselectPreviousEdge();

    // Visually mark this new edge as selected.
    if (selectedEdge) {
      this._lastSelectedEdgeGroup.classed(
          tf.graph.scene.Class.Edge.SELECTED, true);

      // Update the color of the marker too if the edge has one.
      this._updateMarkerOfSelectedEdge(selectedEdge);
    }

    if (this.handleEdgeSelected) {
      // A higher-level component provided a callback. Run it.
      this.handleEdgeSelected(selectedEdge);
    }
  },
  // Called only when a new (non-null) node is selected.
  _nodeSelected: function(event) {
    if (this._allowGraphSelect) {
      this.set('selectedNode', event.detail.name);
    }
    // Reset this variable as a bug in d3 zoom behavior can cause zoomend
    // callback not to be called if a right-click happens during a zoom event.
    this._allowGraphSelect = true;
  },
  _edgeSelected(event) {
    if (this._allowGraphSelect) {
      this.set('_lastSelectedEdgeGroup', event.detail.edgeGroup);
      this.set('selectedEdge', event.detail.edgeData);
    }
    // Reset this variable as a bug in d3 zoom behavior can cause zoomend
    // callback not to be called if a right-click happens during a zoom event.
    this._allowGraphSelect = true;
  },
  _nodeHighlighted: function(event) {
    this.set('highlightedNode', event.detail.name);
  },
  _nodeUnhighlighted: function(event) {
    this.set('highlightedNode', null);
  },
  _nodeToggleExpand: function(event) {
    // Immediately select the node that is about to be expanded.
    this._nodeSelected(event);

    // Compute the sub-hierarchy scene.
    var nodeName = event.detail.name;
    var renderNode = this.renderHierarchy.getRenderNodeByName(nodeName);
    // Op nodes are not expandable.
    if (renderNode.node.type === tf.graph.NodeType.OP) {
      return;
    }
    this.renderHierarchy.buildSubhierarchy(nodeName);
    renderNode.expanded = !renderNode.expanded;

    // Expand the node with some delay so that the user can immediately see
    // the visual effect of selecting that node, before the expansion is
    // done.
    this.async(function() {
      this.querySelector('#scene').setNodeExpanded(renderNode);
    }, 75);
  },
  _nodeToggleExtract: function(event) {
    // Toggle the include setting of the specified node appropriately.
    var nodeName = event.detail.name;
    var renderNode = this.renderHierarchy.getRenderNodeByName(nodeName);
    if (renderNode.node.include == tf.graph.InclusionType.INCLUDE) {
      renderNode.node.include = tf.graph.InclusionType.EXCLUDE;
    } else if (renderNode.node.include == tf.graph.InclusionType.EXCLUDE) {
      renderNode.node.include = tf.graph.InclusionType.INCLUDE;
    } else {
      renderNode.node.include =
       this.renderHierarchy.isNodeAuxiliary(renderNode)
          ? tf.graph.InclusionType.INCLUDE : tf.graph.InclusionType.EXCLUDE;
    }

    // Rebuild the render hierarchy.
    this._buildRenderHierarchy(this.graphHierarchy);
  },
  _nodeToggleSeriesGroup: function(event) {
    // Toggle the group setting of the specified node appropriately.
    var nodeName = event.detail.name;
    tf.graph.toggleNodeSeriesGroup(this.hierarchyParams.seriesMap, nodeName);

    // Rebuild the render hierarchy with the updated series grouping map.
    this.set('progress', {
      value: 0,
      msg: ''
    });
    var tracker = tf.graph.util.getTracker(this);
    var hierarchyTracker = tf.graph.util.getSubtaskTracker(tracker, 100,
          'Namespace hierarchy');
    tf.graph.hierarchy.build(this.basicGraph, this.hierarchyParams, hierarchyTracker)
    .then(function(graphHierarchy) {
      this.set('graphHierarchy', graphHierarchy);
      this._buildRenderHierarchy(this.graphHierarchy);
    }.bind(this));
  },
  _deselectPreviousEdge() {
    const selectedSelector = '.' + tf.graph.scene.Class.Edge.SELECTED;
    // Visually mark the previously selected edge (if any) as deselected.
    d3.select(selectedSelector)
        .classed(tf.graph.scene.Class.Edge.SELECTED, false)
        .each((d, i) => {
          // Reset its marker.
          if (d.label) {
            const paths = d3.select(this).selectAll('path.edgeline');
            if (d.label.startMarkerId) {
              paths.style('marker-start', `url(#${d.label.startMarkerId})`);
            }
            if (d.label.endMarkerId) {
              paths.style('marker-end', `url(#${d.label.endMarkerId})`);
            }
          }
        });
  },
  _updateMarkerOfSelectedEdge(selectedEdge) {
    if (selectedEdge.label) {
      // The marker will vary based on the direction of the edge.
      const markerId = selectedEdge.label.startMarkerId || selectedEdge.label.endMarkerId;
      if (markerId) {
        // Find the corresponding marker for a selected edge.
        const selectedMarkerId = markerId.replace('dataflow-', 'selected-');
        let selectedMarker = this.$$('#' + selectedMarkerId);

        if (!selectedMarker) {
          // The marker for a selected edge of this size does not exist yet. Create it.
          const originalMarker = this.$$('tf-graph-scene').querySelector('#' + markerId);
          selectedMarker = originalMarker.cloneNode(true);
          selectedMarker.setAttribute('id', selectedMarkerId);
          selectedMarker.classList.add('selected-arrowhead');
          originalMarker.parentNode.appendChild(selectedMarker);
        }

        // Make the path use this new marker while it is selected.
        const markerAttribute = selectedEdge.label.startMarkerId ? 'marker-start' : 'marker-end';
        this._lastSelectedEdgeGroup.selectAll('path.edgeline').style(
            markerAttribute, `url(#${selectedMarkerId})`);
      }
    }
  },
  not: function(x) {
    return !x;
  }
});
</script>
